import json
import asyncio
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from openai import OpenAI


class WeatherMCPClient:
    def __init__(self):
        self.session = None
        self.available_tools = []
        self.openai_client = OpenAI(
            base_url="https://api.deepseek.com",
            api_key="sk-129a4b9a8eeb48d3a7648dc81fca0de2",
        )
        self.model = "deepseek-chat"

    async def connect_to_server(self):
        try:
            server_params = StdioServerParameters(
                command="python",
                args=["stdio_server.py"],
                env={"MAP_KEY": "10d653c94cd1009946a54a0e0a8cacab"},
            )
            self._stream_context = stdio_client(server_params)
            stream = await self._stream_context.__aenter__()
            self._session_context = ClientSession(*stream)
            self.session = await self._session_context.__aenter__()
            await self.session.initialize()
            print("MCP客户端初始化成功")
            await self._load_available_tools()

        except Exception as error:
            print(f"连接服务器失败:{str(error)}")
            raise

    async def _load_available_tools(self):
        try:
            tools_response = await self.session.list_tools()
            self.available_tools = tools_response.tools
            print(f"发现了{len(self.available_tools)}个可用工具")
            for tool in self.available_tools:
                print(f"-{tool.name}: {tool.description}")

        except Exception as error:
            print(f"加载工具列表失败:{error}")

    async def cleanup(self):
        try:
            if self._session_context:
                await self._session_context.__aexit__(None, None, None)
            if self._stream_context:
                await self._stream_context.__aexit__(None, None, None)
        except Exception as error:
            print(f"清理资源的时候发生了错误:{str(error)}")

    async def process_request(self, user_input):
        if not self.session:
            return "客户端未连接到服务器"
        try:
            conversion_messages = [
                {
                    "role": "system",
                    "content": "你是天气预报专家，可以调用合适的工具对用户提出的问题进行友好合理的回答",
                },
                {"role": "user", "content": user_input},
            ]
            tool_definitions = []
            tools_response = await self.session.list_tools()
            for tool in tools_response.tools:
                tool_definitions.append(
                    {
                        "type": "function",
                        "name": tool.name,
                        "function": {
                            "name": tool.name,
                            "description": tool.description,
                            "parameters": tool.inputSchema,
                        },
                    }
                )

            first_response = self.openai_client.chat.completions.create(
                model=self.model,
                messages=conversion_messages,
                max_tokens=1024,
                tools=tool_definitions,
            )
            response_texts = []
            first_choice = first_response.choices[0]
            if first_choice.finish_reason == "tool_calls":
                first_tool_call = first_choice.message.tool_calls[0]
                tool_name = first_tool_call.function.name
                tool_arguments = json.loads(first_tool_call.function.arguments)
                conversion_messages.append(
                    {
                        "role": "assistant",
                        "content": None,
                        "tool_calls": [first_tool_call],
                    }
                )
                tool_result = await self.session.call_tool(tool_name, tool_arguments)
                conversion_messages.append(
                    {
                        "role": "tool",
                        "content": tool_result.content[0].text,
                        "tool_call_id": first_tool_call.id,
                    }
                )
                final_response = self.openai_client.chat.completions.create(
                    model=self.model, messages=conversion_messages, max_tokens=1024
                )
                response_texts.append(f"调用工具{tool_name},参数 {tool_arguments}")
                response_texts.append(final_response.choices[0].message.content)
            elif first_choice.finish_reason == "stop":
                response_texts.append(first_choice.message.content)
            return "\n".join(response_texts)
        except Exception as error:
            return f"处理请求时发生了错误:{str(error)}"

    async def chat_loop(self):
        print(f"请输入你的内容，输入quit退出程序")
        while True:
            try:
                user_message = input("\n你的内容:").strip()
                if user_message == "quit":
                    print("感谢使用，再见")
                    break
                if not user_message:
                    print("请输入有效内容")
                    continue
                ai_response = await self.process_request(user_message)
                print("\n" + ai_response)
            except KeyboardInterrupt:
                print("\n程序被用户中断，正在退出...")
                break
            except Exception as error:
                print(f"发生了未预期的错误:{str(error)}")


async def main():
    mcp_client = WeatherMCPClient()
    try:
        # 1.连接MCP服务器
        await mcp_client.connect_to_server()
        # 2.启动智能聊天程序
        await mcp_client.chat_loop()
    except Exception as error:
        print(f"程序运行失败:{str(error)}")
    finally:
        await mcp_client.cleanup()


if __name__ == "__main__":
    asyncio.run(main())
